import tensorflow as tf
from tensorflow import keras
from matplotlib import pyplot as plt
from function_class import UNet_model
from function_class.PP_function import data_only_mix
from function_class.sim_data_gen import original_scene_gen, ordered_image

import tensorflow as tf

# 检查GPU是否可用
if tf.test.is_gpu_available():
    # 如果GPU可用，则设置TensorFlow仅使用GPU
    config = tf.compat.v1.ConfigProto()
    config.gpu_options.allow_growth = True
    session = tf.compat.v1.Session(config=config)

print("GPU加速已启用：", tf.config.list_physical_devices('GPU'))


def ssim_loss(y_true, y_pred):
    ssim = tf.image.ssim(y_true, y_pred, max_val=1.0)
    return 1 - ssim

FOV={'x_len':64,
     'y_len':64,
     'x_detector_position':32,
     'y_detector_position':32}
num_sc=500
md=10
r_d_p = r"E:\pycharm\pycharm_projects\festival\dataset\BW_shadow\row_data"
e_f   = r"E:\pycharm\pycharm_projects\festival\dataset\BW_shadow"

ordered_image(FOV,e_f,r_d_p,file_extension='.png',head_name='img')
x_arr_data,y_tof_data= original_scene_gen(r_d_p,e_f,num_sc,md,FOV,file_extension='.png',head_name='img',
                   new_head_name='scene',new_file_extension='.png',scene_save_file='scene_img',bin=4096)

x_train,x_test,x_valid,y_train,y_test,y_valid=data_only_mix(x_arr_data,y_tof_data)

# 创建学习率调度器
lr_schedule = 0.0001
early_stopping = keras.callbacks.EarlyStopping(
    monitor='val_accuracy',  # 监控验证集准确率（可以改成 'val_loss' 等其他指标）
    patience=10,            # 容忍验证集指标无改善的轮次（这里设置为10轮，可自行调整）
    restore_best_weights=True  # 是否在训练停止后恢复最优权重
)

# 使用包装器
model = UNet_model.BasedOnUNet03(kse=3,kern_reg=None,kern_int_e='he_normal')
adam  = keras.optimizers.Adam(learning_rate=lr_schedule, beta_1=0.9, beta_2=0.999)
model.compile(loss = 'mse',
              optimizer = adam,
              metrics = ['accuracy']
              )
history = model.fit(x_train,y_train,batch_size=10 ,epochs=200,
                    validation_data=(x_valid,y_valid),callbacks=[early_stopping])
model.save(r'E:\pycharm\pycharm_projects\festival\model_save')


# 绘制损失率曲线和准确度曲线
# 提取训练和验证的损失值
train_loss = history.history['loss']
val_loss = history.history['val_loss']
# 提取训练和验证的准确度
train_acc = history.history['accuracy']
val_acc = history.history['val_accuracy']
# 提取 epoch 数量
epochs = range(1, len(train_loss) + 1)
# 创建画布
plt.figure(figsize=(12, 5))

# 绘制损失率曲线
plt.subplot(1, 2, 1)
plt.plot(epochs, train_loss, 'bo-', label='Training Loss')
plt.plot(epochs, val_loss, 'ro-', label='Validation Loss')
plt.title('Training and Validation Loss')
plt.xlabel('Epochs')
plt.ylabel('Loss')
plt.legend()
# 绘制准确度曲线
plt.subplot(1, 2, 2)
plt.plot(epochs, train_acc, 'bo-', label='Training Accuracy')
plt.plot(epochs, val_acc, 'ro-', label='Validation Accuracy')
plt.title('Training and Validation Accuracy')
plt.xlabel('Epochs')
plt.ylabel('Accuracy')
plt.legend()
# 显示图像
plt.tight_layout()
plt.show()
predictions =model.predict(x_test)

for i in range(11):
    plt.figure()
    plt.subplot(121)
    plt.imshow(predictions[i,:,:,0],cmap='viridis_r',vmin=0)
    plt.title('Predicted Image')
    plt.subplot(122)
    plt.imshow(y_test[i,:,:,0],cmap='viridis_r',vmin=0)
    plt.title('ground_truth Image')
    plt.show()

print("测试断点设置处")